# -*- coding: utf-8 -*-
"""
Redis任务管理模块
负责管理列表页任务、详情页任务、失败任务和已完成数据的存储
使用不同Redis DB区分任务类型
"""

import redis
import json
import hashlib
import time
import logging
from typing import Dict, List, Optional, Any
from config.base_settings import REDIS_CONFIG, REDIS_DBS, REDIS_KEYS, get_redis_key

logger = logging.getLogger(__name__)


class RedisManager:
    """Redis任务管理器"""
    
    def __init__(self, spider_name: str):
        """初始化Redis连接"""
        self.spider_name = spider_name
        self.redis_clients = {}
        self._init_redis_clients()
    
    def _get_key(self, key_name: str) -> str:
        """获取带spider_name前缀的Redis key"""
        return get_redis_key(key_name, self.spider_name)
    
    def _init_redis_clients(self):
        """初始化各个Redis客户端"""
        try:
            for db_name, db_num in REDIS_DBS.items():
                self.redis_clients[db_name] = redis.Redis(**REDIS_CONFIG, db=db_num)
                # 测试连接
                self.redis_clients[db_name].ping()
            
            logger.info("Redis连接成功")
        except Exception as e:
            logger.error(f"Redis连接失败: {e}")
            raise
    
    def _get_client(self, db_name: str) -> redis.Redis:
        """获取指定数据库的Redis客户端"""
        return self.redis_clients[db_name]
    
    def generate_md5(self, url: str) -> str:
        """生成URL的MD5哈希值作为唯一键"""
        return hashlib.md5(url.encode('utf-8')).hexdigest()
    
    # 列表页任务管理
    def add_list_task(self, category_url: str, page: int) -> bool:
        """添加列表页任务到队列（去重）"""
        try:
            # 生成任务唯一标识
            task_key = f"{category_url}:{page}"
            
            # 检查是否已存在相同的任务
            if self._is_list_task_exists(category_url, page):
                logger.info(f"列表页任务已存在，跳过: {category_url} 第{page}页")
                return False
            
            task = {
                'category_url': category_url,
                'page': page,
                'status': 'pending',
                'created_at': int(time.time()),
                'retry_count': 0,
                'task_key': task_key
            }
            self._get_client('list_tasks').lpush(self._get_key('LIST_TASKS'), json.dumps(task))
            logger.info(f"添加列表页任务: {category_url} 第{page}页")
            return True
        except Exception as e:
            logger.error(f"添加列表页任务失败: {e}")
            return False
    
    def add_list_task_with_retry(self, category_url: str, page: int, retry_count: int) -> bool:
        """添加带重试次数的列表页任务到队列"""
        try:
            # 生成任务唯一标识
            task_key = f"{category_url}:{page}"
            
            task = {
                'category_url': category_url,
                'page': page,
                'status': 'pending',
                'created_at': int(time.time()),
                'retry_count': retry_count,
                'task_key': task_key
            }
            self._get_client('list_tasks').lpush(self._get_key('LIST_TASKS'), json.dumps(task))
            logger.info(f"添加列表页任务 (重试次数: {retry_count}): {category_url} 第{page}页")
            return True
        except Exception as e:
            logger.error(f"添加列表页任务失败: {e}")
            return False
    
    def _is_list_task_exists(self, category_url: str, page: int) -> bool:
        """检查列表页任务是否已存在"""
        try:
            task_key = f"{category_url}:{page}"
            
            # 检查当前任务队列
            tasks = self._get_client('list_tasks').lrange(self._get_key('LIST_TASKS'), 0, -1)
            for task_data in tasks:
                task = json.loads(task_data)
                if task.get('task_key') == task_key:
                    return True
            
            # 检查失败任务队列
            failed_tasks = self._get_client('list_failed').lrange(self._get_key('LIST_FAILED'), 0, -1)
            for task_data in failed_tasks:
                failed_task = json.loads(task_data)
                task = failed_task.get('task', {})
                if task.get('task_key') == task_key:
                    return True
            
            return False
        except Exception as e:
            logger.error(f"检查列表页任务是否存在失败: {e}")
            return False
    
    def get_list_task_stats(self) -> Dict[str, int]:
        """获取列表页任务统计信息"""
        try:
            # 统计不同类目的任务数量
            tasks = self._get_client('list_tasks').lrange(self._get_key('LIST_TASKS'), 0, -1)
            category_stats = {}
            
            for task_data in tasks:
                task = json.loads(task_data)
                category_url = task.get('category_url', 'unknown')
                category_stats[category_url] = category_stats.get(category_url, 0) + 1
            
            return {
                'total_tasks': len(tasks),
                'category_stats': category_stats
            }
        except Exception as e:
            logger.error(f"获取列表页任务统计失败: {e}")
            return {}
    
    def get_list_task(self) -> Optional[Dict]:
        """获取一个列表页任务（使用原子操作避免竞态条件）"""
        try:
            # 使用原子操作RPOP从队列尾部获取任务
            task_data = self._get_client('list_tasks').rpop(self._get_key('LIST_TASKS'))
            if task_data:
                return json.loads(task_data)
            return None
        except Exception as e:
            logger.error(f"获取列表页任务失败: {e}")
            return None
    
    def get_list_task_count(self) -> int:
        """获取列表页任务数量"""
        try:
            return self._get_client('list_tasks').llen(self._get_key('LIST_TASKS'))
        except Exception as e:
            logger.error(f"获取列表页任务数量失败: {e}")
            return 0
    
    # 详情页任务管理
    def add_detail_task(self, detail_url: str, category_url: str, category_page: int = None) -> bool:
        """添加详情页任务到队列（不去重，允许重复创建）"""
        try:
            task = {
                'detail_url': detail_url,
                'category_url': category_url,
                'category_page': category_page,
                'status': 'pending',
                'created_at': int(time.time()),
                'retry_count': 0
            }
            self._get_client('detail_tasks').lpush(self._get_key('DETAIL_TASKS'), json.dumps(task))
            logger.info(f"添加详情页任务: {detail_url} (类目: {category_url}, 页码: {category_page})")
            return True
        except Exception as e:
            logger.error(f"添加详情页任务失败: {e}")
            return False
    
    def add_detail_task_with_retry(self, detail_url: str, category_url: str, category_page: int, retry_count: int) -> bool:
        """添加带重试次数的详情页任务到队列"""
        try:
            task = {
                'detail_url': detail_url,
                'category_url': category_url,
                'category_page': category_page,
                'status': 'pending',
                'created_at': int(time.time()),
                'retry_count': retry_count
            }
            self._get_client('detail_tasks').lpush(self._get_key('DETAIL_TASKS'), json.dumps(task))
            logger.info(f"添加详情页任务 (重试次数: {retry_count}): {detail_url} (类目: {category_url}, 页码: {category_page})")
            return True
        except Exception as e:
            logger.error(f"添加详情页任务失败: {e}")
            return False
    
    def get_detail_task(self) -> Optional[Dict]:
        """获取一个详情页任务（使用原子操作避免竞态条件）"""
        try:
            # 使用原子操作RPOP从队列尾部获取任务
            task_data = self._get_client('detail_tasks').rpop(self._get_key('DETAIL_TASKS'))
            if task_data:
                return json.loads(task_data)
            return None
        except Exception as e:
            logger.error(f"获取详情页任务失败: {e}")
            return None
    
    def get_detail_task_count(self) -> int:
        """获取详情页任务数量"""
        try:
            return self._get_client('detail_tasks').llen(self._get_key('DETAIL_TASKS'))
        except Exception as e:
            logger.error(f"获取详情页任务数量失败: {e}")
            return 0
    
    # 列表页失败任务管理
    def add_list_failed_task(self, task: Dict, failure_reason: str) -> bool:
        """添加列表页失败任务到失败队列"""
        try:
            failed_task = {
                'task': task,
                'failure_reason': failure_reason,
                'failed_at': int(time.time()),
                'retry_count': task.get('retry_count', 0)  # 任务已经递增了重试次数，不需要再加1
            }
            self._get_client('list_failed').lpush(self._get_key('LIST_FAILED'), json.dumps(failed_task))
            logger.warning(f"添加列表页失败任务: {failure_reason} (重试次数: {failed_task['retry_count']})")
            return True
        except Exception as e:
            logger.error(f"添加列表页失败任务失败: {e}")
            return False
    
    def get_list_failed_task(self) -> Optional[Dict]:
        """获取一个列表页失败任务用于重试"""
        try:
            task_data = self._get_client('list_failed').rpop(self._get_key('LIST_FAILED'))
            if task_data:
                return json.loads(task_data)
            return None
        except Exception as e:
            logger.error(f"获取列表页失败任务失败: {e}")
            return None
    
    def get_list_failed_task_count(self) -> int:
        """获取列表页失败任务数量"""
        try:
            return self._get_client('list_failed').llen(self._get_key('LIST_FAILED'))
        except Exception as e:
            logger.error(f"获取列表页失败任务数量失败: {e}")
            return 0
    
    # 详情页失败任务管理
    def add_detail_failed_task(self, task: Dict, failure_reason: str) -> bool:
        """添加详情页失败任务到失败队列"""
        try:
            failed_task = {
                'task': task,
                'failure_reason': failure_reason,
                'failed_at': int(time.time()),
                'retry_count': task.get('retry_count', 0)  # 任务已经递增了重试次数，不需要再加1
            }
            self._get_client('detail_failed').lpush(self._get_key('DETAIL_FAILED'), json.dumps(failed_task))
            logger.warning(f"添加详情页失败任务: {failure_reason} (重试次数: {failed_task['retry_count']})")
            return True
        except Exception as e:
            logger.error(f"添加详情页失败任务失败: {e}")
            return False
    
    def get_detail_failed_task(self) -> Optional[Dict]:
        """获取一个详情页失败任务用于重试"""
        try:
            task_data = self._get_client('detail_failed').rpop(self._get_key('DETAIL_FAILED'))
            if task_data:
                return json.loads(task_data)
            return None
        except Exception as e:
            logger.error(f"获取详情页失败任务失败: {e}")
            return None
    
    def get_detail_failed_task_count(self) -> int:
        """获取详情页失败任务数量"""
        try:
            return self._get_client('detail_failed').llen(self._get_key('DETAIL_FAILED'))
        except Exception as e:
            logger.error(f"获取详情页失败任务数量失败: {e}")
            return 0
    
    # 已完成数据管理
    def save_completed_data(self, detail_url: str, data: Dict, category_url: str = None, category_page: int = None) -> bool:
        """保存已完成的数据 - 使用URL+类目+页码作为唯一键"""
        try:
            # 使用URL+类目+页码作为唯一键，确保每个类目+页码组合都有独立记录
            unique_key = f"{detail_url}:{category_url}:{category_page}"
            md5_key = self.generate_md5(unique_key)
            
            data_with_key = {
                'detail_url': detail_url,
                'data': data,
                'category_url': category_url,
                'category_page': category_page,
                'completed_at': int(time.time())
            }
            
            # 保存数据
            self._get_client('completed_data').hset(self._get_key('COMPLETED_DATA'), md5_key, json.dumps(data_with_key))
            
            # 维护URL到key的映射（用于快速查找）- 使用专门的DB
            url_mapping_key = f"{self.spider_name}:url_mapping:{detail_url}"
            self._get_client('url_mapping').sadd(url_mapping_key, md5_key)
            
            logger.info(f"保存已完成数据: {detail_url} (类目: {category_url}, 页码: {category_page})")
            return True
        except Exception as e:
            logger.error(f"保存已完成数据失败: {e}")
            return False
    
    def get_completed_data(self, detail_url: str, category_url: str = None, category_page: int = None) -> Optional[Dict]:
        """获取已完成的数据"""
        try:
            if category_url and category_page is not None:
                # 获取特定类目+页码的数据
                unique_key = f"{detail_url}:{category_url}:{category_page}"
                md5_key = self.generate_md5(unique_key)
                data = self._get_client('completed_data').hget(self._get_key('COMPLETED_DATA'), md5_key)
                if data:
                    return json.loads(data)
                return None
            else:
                # 通过URL映射快速查找任意一个已存在的数据记录
                url_mapping_key = f"{self.spider_name}:url_mapping:{detail_url}"
                md5_keys = self._get_client('url_mapping').smembers(url_mapping_key)
                
                if md5_keys:
                    # 获取第一个找到的数据记录
                    for md5_key in md5_keys:
                        data = self._get_client('completed_data').hget(self._get_key('COMPLETED_DATA'), md5_key)
                        if data:
                            return json.loads(data)
                return None
        except Exception as e:
            logger.error(f"获取已完成数据失败: {e}")
            return None
    
    def is_url_scraped(self, detail_url: str) -> bool:
        """检查URL是否已经被抓取过（用于去重）"""
        try:
            # 通过URL映射快速检查
            url_mapping_key = f"{self.spider_name}:url_mapping:{detail_url}"
            return self._get_client('url_mapping').exists(url_mapping_key)
        except Exception as e:
            logger.error(f"检查URL是否已抓取失败: {e}")
            return False
    
    def update_completed_data_category(self, detail_url: str, category_url: str, category_page: int = None) -> bool:
        """更新已完成数据的类目URL和页码"""
        try:
            md5_key = self.generate_md5(detail_url)
            data = self._get_client('completed_data').hget(self._get_key('COMPLETED_DATA'), md5_key)
            if data:
                data_dict = json.loads(data)
                data_dict['category_url'] = category_url
                if category_page is not None:
                    data_dict['category_page'] = category_page
                self._get_client('completed_data').hset(self._get_key('COMPLETED_DATA'), md5_key, json.dumps(data_dict))
                logger.info(f"更新数据类目: {detail_url} (类目: {category_url}, 页码: {category_page})")
                return True
            return False
        except Exception as e:
            logger.error(f"更新数据类目失败: {e}")
            return False
    
    def is_data_completed(self, detail_url: str, category_url: str = None, category_page: int = None) -> bool:
        """检查数据是否已完成"""
        try:
            if category_url and category_page is not None:
                # 检查特定类目+页码的数据是否存在
                unique_key = f"{detail_url}:{category_url}:{category_page}"
                md5_key = self.generate_md5(unique_key)
                return self._get_client('completed_data').hexists(self._get_key('COMPLETED_DATA'), md5_key)
            else:
                # 兼容旧方法，只基于URL检查
                md5_key = self.generate_md5(detail_url)
                return self._get_client('completed_data').hexists(self._get_key('COMPLETED_DATA'), md5_key)
        except Exception as e:
            logger.error(f"检查数据完成状态失败: {e}")
            return False
    
    def get_all_completed_data(self) -> List[Dict]:
        """获取所有已完成的数据"""
        try:
            data_dict = self._get_client('completed_data').hgetall(self._get_key('COMPLETED_DATA'))
            result = []
            for key, value in data_dict.items():
                try:
                    data = json.loads(value)
                    result.append(data)
                except json.JSONDecodeError:
                    continue
            return result
        except Exception as e:
            logger.error(f"获取所有已完成数据失败: {e}")
            return []
    
    # 任务锁管理
    def acquire_task_lock(self, lock_key: str, timeout: int = 30) -> bool:
        """获取任务锁"""
        try:
            return self._get_client('list_tasks').set(lock_key, "locked", nx=True, ex=timeout)
        except Exception as e:
            logger.error(f"获取任务锁失败: {e}")
            return False
    
    def release_task_lock(self, lock_key: str) -> bool:
        """释放任务锁"""
        try:
            return self._get_client('list_tasks').delete(lock_key) > 0
        except Exception as e:
            logger.error(f"释放任务锁失败: {e}")
            return False
    
    # 批量操作
    def batch_add_detail_tasks(self, tasks: List[Dict]) -> int:
        """批量添加详情页任务"""
        try:
            if not tasks:
                return 0
            
            pipe = self._get_client('detail_tasks').pipeline()
            for task in tasks:
                pipe.lpush(self._get_key('DETAIL_TASKS'), json.dumps(task))
            results = pipe.execute()
            logger.info(f"批量添加{len(tasks)}个详情页任务")
            return len(tasks)
        except Exception as e:
            logger.error(f"批量添加详情页任务失败: {e}")
            return 0
    
    def batch_add_list_tasks_fast(self, tasks: List[Dict]) -> int:
        """快速批量添加列表页任务（跳过去重检查）"""
        try:
            if not tasks:
                return 0
            
            pipe = self._get_client('list_tasks').pipeline()
            for task in tasks:
                pipe.lpush(self._get_key('LIST_TASKS'), json.dumps(task))
            results = pipe.execute()
            logger.info(f"快速批量添加{len(tasks)}个列表页任务")
            return len(tasks)
        except Exception as e:
            logger.error(f"快速批量添加列表页任务失败: {e}")
            return 0
    
    # 统计信息
    def get_stats(self) -> Dict[str, int]:
        """获取任务统计信息"""
        try:
            return {
                'list_tasks': self.get_list_task_count(),
                'detail_tasks': self.get_detail_task_count(),
                'list_failed_tasks': self.get_list_failed_task_count(),
                'detail_failed_tasks': self.get_detail_failed_task_count(),
                'completed_data': self._get_client('completed_data').hlen(self._get_key('COMPLETED_DATA'))
            }
        except Exception as e:
            logger.error(f"获取统计信息失败: {e}")
            return {}
    
    def clear_all_tasks(self) -> bool:
        """清空所有任务（谨慎使用）"""
        try:
            for db_name in ['list_tasks', 'detail_tasks', 'list_failed', 'detail_failed']:
                self._get_client(db_name).delete(self._get_key('LIST_TASKS'))
                self._get_client(db_name).delete(self._get_key('DETAIL_TASKS'))
                self._get_client(db_name).delete(self._get_key('LIST_FAILED'))
                self._get_client(db_name).delete(self._get_key('DETAIL_FAILED'))
            logger.warning("已清空所有任务")
            return True
        except Exception as e:
            logger.error(f"清空任务失败: {e}")
            return False
    
    def clear_completed_data(self) -> bool:
        """清空已完成数据（谨慎使用）"""
        try:
            self._get_client('completed_data').delete(self._get_key('COMPLETED_DATA'))
            logger.warning("已清空已完成数据")
            return True
        except Exception as e:
            logger.error(f"清空已完成数据失败: {e}")
            return False
    
    def get_all_list_tasks(self) -> List[Dict]:
        """获取所有列表任务"""
        try:
            tasks = []
            for task_data in self._get_client('list_tasks').lrange(self._get_key('LIST_TASKS'), 0, -1):
                tasks.append(json.loads(task_data))
            return tasks
        except Exception as e:
            logger.error(f"获取所有列表任务失败: {e}")
            return []
    
    def get_all_detail_tasks(self) -> List[Dict]:
        """获取所有详情任务"""
        try:
            tasks = []
            for task_data in self._get_client('detail_tasks').lrange(self._get_key('DETAIL_TASKS'), 0, -1):
                tasks.append(json.loads(task_data))
            return tasks
        except Exception as e:
            logger.error(f"获取所有详情任务失败: {e}")
            return []
    
    def get_all_list_failed_tasks(self) -> List[Dict]:
        """获取所有列表失败任务"""
        try:
            tasks = []
            for task_data in self._get_client('list_failed').lrange(self._get_key('LIST_FAILED'), 0, -1):
                tasks.append(json.loads(task_data))
            return tasks
        except Exception as e:
            logger.error(f"获取所有列表失败任务失败: {e}")
            return []
    
    def get_all_detail_failed_tasks(self) -> List[Dict]:
        """获取所有详情失败任务"""
        try:
            tasks = []
            for task_data in self._get_client('detail_failed').lrange(self._get_key('DETAIL_FAILED'), 0, -1):
                tasks.append(json.loads(task_data))
            return tasks
        except Exception as e:
            logger.error(f"获取所有详情失败任务失败: {e}")
            return []
    
    def get_all_url_mapping(self) -> Dict:
        """获取所有URL映射"""
        try:
            return self._get_client('url_mapping').hgetall(self._get_key('URL_MAPPING'))
        except Exception as e:
            logger.error(f"获取所有URL映射失败: {e}")
            return {}
    
    def clear_url_mapping(self) -> bool:
        """清空URL映射"""
        try:
            # 清空所有url_mapping:*键
            client = self._get_client('url_mapping')
            keys = client.keys(f'{self.spider_name}:url_mapping:*')
            if keys:
                client.delete(*keys)
            logger.info(f"URL映射已清空，删除了 {len(keys)} 个映射键")
            return True
        except Exception as e:
            logger.error(f"清空URL映射失败: {e}")
            return False
    
    # 统计信息管理
    def set_stats(self, stats_data: Dict) -> bool:
        """设置统计信息"""
        try:
            self._get_client('stats').hset(self._get_key('STATS'), mapping=stats_data)
            return True
        except Exception as e:
            logger.error(f"设置统计信息失败: {e}")
            return False
    
    def get_stats_data(self) -> Dict:
        """获取统计信息"""
        try:
            return self._get_client('stats').hgetall(self._get_key('STATS'))
        except Exception as e:
            logger.error(f"获取统计信息失败: {e}")
            return {}
    
    def update_stats(self, key: str, value: int) -> bool:
        """更新统计信息"""
        try:
            self._get_client('stats').hset(self._get_key('STATS'), key, value)
            return True
        except Exception as e:
            logger.error(f"更新统计信息失败: {e}")
            return False
    
    def increment_stats(self, key: str, increment: int = 1) -> bool:
        """增加统计信息"""
        try:
            self._get_client('stats').hincrby(self._get_key('STATS'), key, increment)
            return True
        except Exception as e:
            logger.error(f"增加统计信息失败: {e}")
            return False
    
    def clear_stats(self) -> bool:
        """清空统计信息"""
        try:
            self._get_client('stats').delete(self._get_key('STATS'))
            logger.info("统计信息已清空")
            return True
        except Exception as e:
            logger.error(f"清空统计信息失败: {e}")
            return False
    
    # Wayfair 搜索词任务管理
    def add_keyword_task(self, task: Dict) -> bool:
        """添加搜索词任务"""
        try:
            keyword_key = f"{self.spider_name}:keyword_tasks"
            self._get_client('list_tasks').lpush(keyword_key, json.dumps(task))
            logger.info(f"添加搜索词任务: {task.get('keyword', 'N/A')}")
            return True
        except Exception as e:
            logger.error(f"添加搜索词任务失败: {e}")
            return False
    
    def get_keyword_task(self) -> Optional[Dict]:
        """获取一个搜索词任务"""
        try:
            keyword_key = f"{self.spider_name}:keyword_tasks"
            task_data = self._get_client('list_tasks').rpop(keyword_key)
            if task_data:
                return json.loads(task_data)
            return None
        except Exception as e:
            logger.error(f"获取搜索词任务失败: {e}")
            return None
    
    def get_keyword_task_count(self) -> int:
        """获取搜索词任务数量"""
        try:
            keyword_key = f"{self.spider_name}:keyword_tasks"
            return self._get_client('list_tasks').llen(keyword_key)
        except Exception as e:
            logger.error(f"获取搜索词任务数量失败: {e}")
            return 0
    
    def add_keyword_failed_task(self, task: Dict, failure_reason: str) -> bool:
        """添加失败的搜索词任务"""
        try:
            failed_task = {
                'task': task,
                'failure_reason': failure_reason,
                'failed_at': int(time.time())
            }
            keyword_failed_key = f"{self.spider_name}:keyword_failed"
            self._get_client('list_failed').lpush(keyword_failed_key, json.dumps(failed_task))
            logger.warning(f"添加搜索词失败任务: {failure_reason}")
            return True
        except Exception as e:
            logger.error(f"添加搜索词失败任务失败: {e}")
            return False
    
    def save_keyword_data(self, keyword: str, data: Dict) -> bool:
        """保存搜索词数据"""
        try:
            keyword_data_key = f"{self.spider_name}:keyword_data"
            md5_key = self.generate_md5(keyword)
            
            data_with_meta = {
                'keyword': keyword,
                'data': data,
                'saved_at': int(time.time())
            }
            
            self._get_client('completed_data').hset(keyword_data_key, md5_key, json.dumps(data_with_meta))
            logger.info(f"保存搜索词数据: {keyword}")
            return True
        except Exception as e:
            logger.error(f"保存搜索词数据失败: {e}")
            return False
    
    def get_all_keyword_data(self) -> List[Dict]:
        """获取所有搜索词数据"""
        try:
            keyword_data_key = f"{self.spider_name}:keyword_data"
            data_dict = self._get_client('completed_data').hgetall(keyword_data_key)
            result = []
            for key, value in data_dict.items():
                try:
                    data = json.loads(value)
                    result.append(data)
                except json.JSONDecodeError:
                    continue
            return result
        except Exception as e:
            logger.error(f"获取所有搜索词数据失败: {e}")
            return []